package cn.jbolt.ai.klb.service.impl;

import cn.hutool.core.date.DateTime;
import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.core.fileSplit.FileSplitCache;
import cn.jbolt.ai.core.vdb.VdbCollectionService;
import cn.jbolt.ai.klb.DocLinkType;
import cn.jbolt.ai.klb.FileType;
import cn.jbolt.ai.klb.entity.AiDocAnswer;
import cn.jbolt.ai.klb.service.AiDocAnswerService;
import cn.jbolt.util.JSONUtil;
import cn.jbolt.util.Pager;
import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import cn.jbolt.ai.klb.DocDataType;
import cn.jbolt.ai.klb.FileState;
import cn.jbolt.ai.klb.dao.AiDocMapper;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.ai.klb.entity.AiFile;
import cn.jbolt.ai.klb.entity.AiKnowledgeBase;
import cn.jbolt.ai.klb.service.AiDocService;
import cn.jbolt.ai.klb.service.AiFileService;
import cn.jbolt.ai.klb.service.AiKnowledgeBaseService;
import cn.jbolt.util.Result;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.model.JBoltAI;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.resource.ai.AIModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;


@Service
public class AiDocServiceImpl extends ServiceImpl<AiDocMapper, AiDoc> implements AiDocService {

    private static final Logger logger = LoggerFactory.getLogger(AiDocServiceImpl.class);
    @Autowired
    private AiKnowledgeBaseService aiKnowledgeBaseService;
    @Autowired
    private VdbCollectionService vdbCollectionService;
    @Autowired
    private AiFileService aiFileService;
    @Autowired
    private AiDocMapper aiDocMapper;

    @Autowired
    protected AiDocAnswerService docAnswerService;

    @Autowired
    @Qualifier("threadPoolTaskExecutor")
    private ThreadPoolTaskExecutor threadPool;

    @Override
    public Result<AiDoc> getDocById(Long id) {
        AiDoc doc = aiDocMapper.getDocById(id);
        if (doc != null) {

            return Result.success(doc);
        } else {
            return Result.error("文档不存在");
        }
    }

    @Override
    public Result listByPage(Integer pageNumber, Integer pageSize, String keywords, Long aiKnowledgeBaseId, String dataType,
                             Integer status) {
        List<AiDoc> docFile = aiDocMapper.getList((pageNumber - 1) * pageSize, pageSize, keywords,
                aiKnowledgeBaseId, dataType, status);
        Page page = new Page(pageNumber, pageSize, aiDocMapper.getListCount(keywords, aiKnowledgeBaseId, dataType, status));
        page.setRecords(docFile);
        return Result.success(Pager.from(page));
    }

    @Override
    @Transactional
    public Result add(AiDoc doc) {
        Result result = aiKnowledgeBaseService.checkKnowledgeBase(doc.getKnowledgeBaseId());
        if (result.isError()) {
            return result;
        }
        AiKnowledgeBase knowledgeBase = (AiKnowledgeBase) result.getData();
        //判断
        if (doc.getDocType().equals(FileType.QA.getCode())) {  //QA.把答案添加到answer表
            AiDocAnswer docAnswer = new AiDocAnswer();
            docAnswer.setAnswer(doc.getAnswer());
            Result<AiDocAnswer> addResult = docAnswerService.add(docAnswer);
            if (addResult.isError()) {
                return addResult;
            }
            docAnswer = addResult.getData();
            doc.setLinkId(docAnswer.getId());
            doc.setLinkType(DocLinkType.LINK_TYPE_ANSWER.getCode());
        }
        doc.setDocStatus(FileState.TRAIN_SUCCESS.getCode());//只要是有了doc对象,肯定是训练成功的对象
        doc.setTrainTime(DateTime.now());
        boolean save = this.save(doc);
        if (save) {
            //异步向量化doc
            threadPool.submit(() -> embeddingDoc(doc, doc.getContent(), knowledgeBase.getVdbResourceId().toString(), doc.getKnowledgeBaseId()));
            return Result.success();
        } else {
            return Result.error("添加失败");
        }
    }

    @Override
    public Result addQuestionByDocId(AiDoc doc) {
        if (doc == null) {
            return Result.error("参数错误");
        }
        doc.setDocStatus(FileState.TRAIN_SUCCESS.getCode());
        //  取出一个知识库的iD,查询vdbId
        Result result = aiKnowledgeBaseService.checkKnowledgeBase(doc.getKnowledgeBaseId());
        if (result.isError()) {
            return result;
        }
        AiKnowledgeBase knowledgeBase = (AiKnowledgeBase) result.getData();
        boolean save = this.save(doc);
        if (save) {
            //异步向量化doc
            threadPool.submit(() -> embeddingDoc(doc, doc.getContent(), knowledgeBase.getVdbResourceId().toString(), doc.getKnowledgeBaseId()));

            return Result.success();
        } else {
            return Result.error("添加失败");
        }

    }

    @Override
    public Result getQuestionListByLinkId(Long docId) {
        if (docId == null) {
            return Result.error("查询失败");
        }
        QueryWrapper<AiDoc> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("doc_id", docId);
        List<AiDoc> docList = this.list(queryWrapper);
        return Result.success(docList);
    }

    private static final String AI_GENERATE_Q_TPL = "现在是一个问题生成专家。你的唯一任务是从输入文本中提取关键信息，并精确生成2-3个高质量问题。\n" +
            "必须遵循以下严格规则：\n" +
            "\n" +
            "问题必须采用普通人日常使用的口语化表达方式\n" +
            "问题必须与文本内容高度相关，不允许偏离主题\n" +
            "问题必须包含文本中的关键词和核心概念\n" +
            "问题应覆盖文本中最重要的信息点和普通读者最可能关注的方面\n" +
            "问题应具有多样性，避免重复相同的信息点\n" +
            "\n" +
            "输出格式要求：\n" +
            "\n" +
            "必须且只能输出一个有效的JSON数组格式\n" +
            "格式必须严格为：[\"问题1\", \"问题2\", \"问题3\"]\n" +
            "不允许输出任何其他内容，包括说明、解释或注释\n" +
            "不允许使用markdown代码块或其他格式化标记\n" +
            "不允许在JSON前后添加任何文字\n" +
            "确保JSON格式完全正确，可以被直接解析\n\n" +
            "文本如下：\n{}";

    @Transactional
    @Override
    public Result generateQuestionByModel(String modelSn, String content, AiDoc doc) {
        if (StrUtil.isBlank(modelSn)) {
            return Result.error("请选择大模型");
        }
        if (StrUtil.isBlank(content)) {
            return Result.error("内容不能为空");
        }
        String[] model = modelSn.split(":", 2);
        AIChatEvent event = JBoltAI.chat();
        event.setModel(AIModel.get(model[0]));
        if (model.length > 1 && StrUtil.isNotEmpty(model[1])) {
            event.setModelName(model[1]);
        }
        event.setAttr("format", "json");
        event.setStream(false);
        AIMessage res = event.prompt(AI_GENERATE_Q_TPL, content).get();
        if (event.isFail()) {
            return Result.error(event.getError().getMsg());
        }
        try {

            ArrayNode questions = JSONUtil.toJsonArray(res.getContent());
            for (int i = 0; i < questions.size(); i++) {
                String c = questions.get(i).asText();
                doc.setContent(c);
                doc.setId(null);
                addQuestionByDocId(doc);
            }
            return Result.success(questions);
        } catch (Exception e) {
            e.printStackTrace();
            return Result.error("AI生成内容失败，请重新尝试");
        }

    }

    public void embeddingDoc(AiDoc aiDoc, String text, String vdbResourceId, Long knowledgeBaseId) {
        try {
            boolean fileHasFail = FileSplitCache.hasFailed(aiDoc.getFileId());
            if (fileHasFail) {
                return;
            }
            //向量化中的状态
            aiDoc.setDocStatus(FileState.EMBEDDING.getCode());
            this.updateById(aiDoc);
            if (StrUtil.isNotEmpty(text)) {
                Result result = vdbCollectionService.insertData(vdbResourceId,
                        "KDB_" + knowledgeBaseId, aiDoc.getId().toString(), text,
                        "file_id", (aiDoc.getFileId() != null ? aiDoc.getFileId().toString() : "")
                        , "data_type", aiDoc.getDataType());
                if (result.isError()) {
                    //更新doc中的表的状态
                    aiDoc.setDocStatus(FileState.EMBEDDING_ERROR.getCode());//向量化失败
                    this.updateById(aiDoc);
                    FileSplitCache.addFailCount(aiDoc.getFileId());
                } else {
                    //更新doc中的表的状态
                    aiDoc.setDocStatus(FileState.EMBEDDING_SUCCESS.getCode());//向量化成功
                    aiDoc.setEmbedTime(DateTime.now());
                    this.updateById(aiDoc);
                    FileSplitCache.addSuccessCount(aiDoc.getFileId());
                }
            }

        } catch (Exception e) {
            e.printStackTrace();
            logger.error("向量化异常：" + JSON.toJSONString(aiDoc));

        }
    }

    @Override
    @Transactional
    public Result update(AiDoc doc) {
        if (StrUtil.isEmptyIfStr(doc.getContent())) {
            return Result.error("内容不能为空");
        }
        if (updateById(doc)) {

            AiKnowledgeBase aiKnowledgeBase = aiKnowledgeBaseService.getById(doc.getKnowledgeBaseId());
            if (aiKnowledgeBaseService == null) {
                return Result.error("知识库不存在");
            }
            //先移除VDB中的数据
            Result removeResult = vdbCollectionService.removeDataById(aiKnowledgeBase.getVdbResourceId().toString(),
                    "KDB_" + doc.getKnowledgeBaseId(), doc.getId().toString());
            if (removeResult.isError()) {
                return removeResult;
            }

            //在添加VDB中的数据
            Result result = vdbCollectionService.insertData(aiKnowledgeBase.getVdbResourceId().toString(),
                    "KDB_" + doc.getKnowledgeBaseId(),
                    doc.getId().toString(), doc.getContent(),
                    "file_id", doc.getId().toString()
                    , "data_type", DocDataType.FILE_PIECE.getCode());
            if (result.isError()) {
                //更新doc中的表的状态
                doc.setDocStatus(FileState.EMBEDDING_ERROR.getCode());//向量化失败
                updateById(doc);
                return result;
            } else {
                //更新doc中的表的状态
                doc.setDocStatus(FileState.EMBEDDING_SUCCESS.getCode());//向量化成功
                doc.setEmbedTime(new Date());
                updateById(doc);
            }

            return Result.success();
        } else {
            return Result.error("更新失败");
        }
    }

    @Override
    @Transactional
    public Result delete(Long id) {
        Result<AiDoc> fileResult = getDocById(id);
        if (fileResult.isError()) return fileResult;
        AiDoc doc = fileResult.getData();
        if (doc.getLinkType() != null && doc.getDocId() == null) {
            if (doc.getLinkType().equals(DocLinkType.LINK_TYPE_ANSWER.getCode())) { //如果是QA的话要移除他对应的answeer
                Result<AiDocAnswer> result = docAnswerService.delete(doc.getLinkId());
                if (result.isError()) return result;
            }
        }
        if (removeById(id)) {
            //移除相应的vdb
            Result knowledgeBaseResult = aiKnowledgeBaseService.getKnowledgeBaseById(doc.getKnowledgeBaseId());
            if (knowledgeBaseResult.isError()) return knowledgeBaseResult;
            AiKnowledgeBase knowledgeBase = (AiKnowledgeBase) knowledgeBaseResult.getData();
            Result removeVDBResult = vdbCollectionService.removeDataById(knowledgeBase.getVdbResourceId().toString(),
                    "KDB_" + knowledgeBase.getId(), doc.getId().toString());
            if (removeVDBResult.isError()) return Result.error("移除vdb失败");
            //判断是否有基于二次训练的数据
            Result result = checkTwiceTrain(id);
            if (result.isError()) return result;
            return Result.success();
        } else {
            return Result.error("删除失败");
        }
    }

    /**
     * 判断当前数据是否有基于二次训练的数据
     *
     * @param id
     * @return
     */
    private Result checkTwiceTrain(Long id) {
        QueryWrapper<AiDoc> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("doc_id", id);
        List<AiDoc> list = this.list(queryWrapper);
        if (list.size() > 0) {
            for (AiDoc doc : list) {
                Result result = delete(doc.getId());
                if (result.isError()) {
                    return result;
                }
            }

        }
        return Result.success();
    }


    public AiDoc getKnowledgeBaseName(AiDoc aiDoc) {
        if (aiDoc.getKnowledgeBaseId() != null) {
            Result result = aiKnowledgeBaseService.getKnowledgeBaseById(aiDoc.getKnowledgeBaseId());
            if (result.isSuccess()) {
                AiKnowledgeBase aiKnowledgeBase = (AiKnowledgeBase) result.getData();
                aiDoc.setAiKnowledgeBaseName(aiKnowledgeBase.getName());
            }
        }
        return aiDoc;
    }


    @Override
    public Result getDataTypeList() {
        return Result.success(DocDataType.toList());
    }

    /**
     * 根据知识库id移除DOC
     *
     * @return
     */
    @Override
    public Result delDocByKlbId(Long klbId) {
        QueryWrapper<AiDoc> docQuery = new QueryWrapper<AiDoc>().eq("knowledge_base_id", klbId);
        List<AiDoc> docList = this.list(docQuery);
        if (docList != null && docList.size() > 0) {
            boolean remove = this.remove(docQuery);
            if (remove == false) {
                return Result.error("删除失败");
            }
        }
        return Result.success();
    }

    /**
     * 根据文件id移除DOC
     *
     * @param fileId
     * @return
     */
    @Override
    public Result delDocByFileId(Long fileId) {
        QueryWrapper<AiDoc> docQuery = new QueryWrapper<AiDoc>().eq("file_id", fileId).eq("data_type", DocDataType.FILE_PIECE.getCode());
        List<AiDoc> docList = this.list(docQuery);
        if (docList != null && docList.size() > 0) {
            boolean remove = this.remove(docQuery);
            if (remove == false) {
                return Result.error("删除失败");
            }
        }
        return Result.success();
    }


    @Override
    public List<AiDoc> getDocListByIds(List<String> ids) {
        if (ids != null && ids.size() > 0) {
            return aiDocMapper.getDocListByIds(ids);
        }
        return Collections.emptyList();
    }

    @Override
    public Result<List<AiDoc>> getNormalDocListByFileId(Long fileId) {
        QueryWrapper<AiDoc> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("file_id", fileId);
        queryWrapper.isNull("doc_id");
        if (this.list(queryWrapper) != null) {
            return Result.success(this.list(queryWrapper));
        }
        return Result.success(new ArrayList<>());
    }

    @Override
    public Result<List<AiDoc>> getQADocListByFileId(Long fileId) {
        return Result.success(aiDocMapper.getQADocListByFileId(fileId));

    }
}

